import cv2
import torch
from pathlib import Path

# 加载 YOLOv5 模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# 打开视频文件（如果使用摄像头，则可以将参数设为 0 或其他摄像头索引）
video_path = 'b.mp4'  # 替换为你的视频文件路径
cap = cv2.VideoCapture(video_path)

while cap.isOpened():
    ret, frame = cap.read()

    if not ret:
        break

    # 将视频帧调整为模型期望的输入尺寸
    img = cv2.resize(frame, (640, 480))  # 将尺寸调整为模型期望的大小，可以根据模型配置进行调整

    # 将 OpenCV 图像转换为 PyTorch 张量
    img = torch.from_numpy(img).permute(2, 0, 1).float().div(255.0).unsqueeze(0)

    # 进行目标检测
    results = model(img)

    # 获取检测结果
    pred = results[0]

    # 将检测结果绘制到帧上
    for det in pred:
        if det[-1] > 0.5:  # 这里设置信度阈值为0.5
            x1, y1, x2, y2, conf, cls = det[:6]
            x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
            label = model.names[int(cls)]
            color = (0, 255, 0)  # 框的颜色，绿色
            frame = cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            frame = cv2.putText(frame, f'{label}: {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    # 显示帧
    cv2.imshow('YOLOv5 Video Detection', frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()
